from langchain.prompts import PromptTemplate
from LLM_Prompt.siliconflow import SiliconFlow

# 使用 SiliconFlow 或 OpenAI 接入大模型
llm = SiliconFlow()  # 替换成你实际的 SiliconFlow 模型接口

# 问题分类 Prompt 模板
classification_prompt = PromptTemplate(
    input_variables=["question"],
    # template="请将以下问题的内容分类为以下几类之一：\n问题: {question} \n- 轨迹查询 \n- 图片调度 \n- 天气查询 \n- 其他 \n只需回答上面列出的一个分类名称，不要添加任何解释。",
    template="请将以下问题的内容分类为以下几类之一：\n问题: {question}\n- 船舶流量\n- 过船记录\n- 统计问题\n- 天气查询 \n- 其他\n只需回答上面列出的一个分类名称，不要添加任何解释。",
)

classification_chain = classification_prompt | llm

def classify(question):
    # 1. 问题分类
    classification_result = classification_chain.invoke({"question": question})
    # print("分类结果: ", classification_result)

    return classification_result

# example_question = "北京2025年1月10日天气怎么样"
# result = classify(example_question)
# print(result)